import unittest
import sys
sys.path.append("..")

from psychsim.pwl import *
from psychsim.action import Action,ActionSet
from psychsim.world import World,stateKey
from psychsim.agent import Agent
from psychsim.reward import *

''' 
This example is used to test setting an agent's belief on its own models:
To what extend are these beliefs considered for the action selection?
'''

class AmILazyTestCase(unittest.TestCase):

    def setUp(self):
        #intialize world, agent, order, and max rounds
        self.world = World()
        self.sammy = Agent('Sammy')
        self.world.addAgent(self.sammy)
        
        self.world.setOrder(['Sammy'])
        self.maxRounds=5

    def addStates(self):
        #add states for round and energy
        self.world.defineState(None,'round',int)
        self.world.setState(None,'round',0)
        
        self.world.defineState('Sammy','energy',int)
        self.world.setState('Sammy','energy',10)
                    
    def addTermination(self):
        #terminate if maxRounds reached
        self.world.addTermination(makeTree({'if': thresholdRow(stateKey(None,'round'),self.maxRounds),
                                   True: True, False: False}))

    def addActions(self):
        #specify Sammy's actions and dynamics
        action =  self.sammy.addAction({'verb': 'do nothing'})
        tree = makeTree(incrementMatrix(stateKey(None,'round'),1))
        self.world.setDynamics('round',action,tree) 
        
        action =  self.sammy.addAction({'verb': 'do something'})
        tree = makeTree(incrementMatrix(stateKey('Sammy','energy'),-1))
        self.world.setDynamics(stateKey('Sammy','energy'),action,tree) 
        tree = makeTree(incrementMatrix(stateKey(None,'round'),1))
        self.world.setDynamics('round',action,tree) 

    def addModels(self):
        #add models for Sammy's belief that he is lazy or energetic
        self.sammy.addModel('Lazy',R={},rationality=1.,level=2)
        self.sammy.addModel('Energetic',R={},rationality=1.,level=2)
        
        beliefs = {'Lazy':0.9, 'Energetic':0.1}
         
        self.sammy.world.setMentalModel('Sammy','Sammy', beliefs)
        self.world.setModel('Sammy', True)

    def setRewards(self):
        #add rewards
        self.sammy.setReward(minimizeFeature(stateKey('Sammy','energy')), 5., True)
        self.sammy.setReward(maximizeFeature(stateKey('Sammy','energy')), 1., 'Lazy')
        self.sammy.setReward(minimizeFeature(stateKey('Sammy','energy')), 1., 'Energetic')
                                                                        
    def testSetup(self):
        #setup everything
        self.addStates()
        self.addTermination()
        self.addActions()
        self.addModels()
        self.setRewards()
            
    def testRun(self):
        self.testSetup()
        self.world.printState()
        #self.world.printModels()
        for t in range(self.maxRounds + 1):
            self.world.explain(self.world.step(),level=1)
            #self.world.state.select()
            self.world.printState()
            #self.world.printModels()
            if self.world.terminated():
                break 
            

if __name__ == "__main__":
    suite = unittest.TestSuite()
    for method in dir(AmILazyTestCase):
       if method.startswith("test"):
          suite.addTest(AmILazyTestCase(method))
    unittest.TextTestRunner().run(suite)
            
